import pandas as pd
import sklearn.preprocessing
import joblib
from model.Mymodel import Mymodel
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer, MinMaxScaler, StandardScaler
import os


def loadData():
    data_list = []
    path = 'D:\\file\\DS\\yearaqi'
    cna = ['city', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3', 'CO', 'next_level']
    alpath = os.listdir(path)[5:]
    for file in alpath:
        tmp = pd.read_csv(path + '//' + file)[cna]
        data_list.append(tmp)
    all_data = pd.concat(data_list)
    target_var = 'next_level'  # 目标变量
    # 数据集特征
    features = list(all_data.columns)
    features.remove(target_var)
    # 目标变量的类别字典
    Class = all_data[target_var].unique()
    Class_dic = dict(zip(Class, range(len(Class))))
    all_data['target'] = all_data[target_var].apply(lambda x: Class_dic[x])
    # 对目标变量进行0-1编码
    lb = LabelBinarizer()
    lb.fit(list(Class_dic.values()))
    transformed_lables = lb.transform(all_data['target'])
    y_bin_labels = []  # 对多分类进行0-1编码的变量
    for i in range(transformed_lables.shape[1]):
        y_bin_labels.append('y' + str(i))
        all_data['y' + str(i)] = transformed_lables[:, i]
    # scaler = StandardScaler()
    scaler = MinMaxScaler()
    all_data[features] = scaler.fit_transform(all_data[features])
    joblib.dump(scaler, 'scaler.pkl')
    train_x, test_x, train_y, test_y = train_test_split(all_data[features],
                                                        all_data[y_bin_labels],
                                                        train_size=0.7,
                                                        test_size=0.3,
                                                        random_state=12)

    return train_x, train_y, test_x, test_y


model = Mymodel(7, 7, 6)
scaler = StandardScaler()
mini = MinMaxScaler()
x_train, y_train, x_test, y_test = loadData()
# x_train=scaler.fit_transform(x_train)
# x_train = mini.fit_transform(x_train)
model.train(x_train, y_train, lr=0.001, epochs=200)
# x_test=scaler.fit_transform(x_test)
# x_test = mini.fit_transform(x_test)
model.test(x_test, y_test)
model.save()
